import os
from mistralai import Mistral

# Initialize the client
# api_key = os.environ.get("MISTRAL_API_KEY")
api_key = "Kz2yZgyEQxkreQRagET0rtz15l8smbRT"

client = Mistral(api_key)

# 1. 初始化一个列表来存储对话历史
conversation_history = []

while True:
    # 2. 获取用户输入
    user_input = input("你: ")
    if user_input.lower() in ["exit", "quit"]:
        break

    # 3. 将用户输入添加到对话历史中
    conversation_history.append({"role": "user", "content": user_input})

    try:
        # 4. 发送完整的对话历史到 API
        response = client.beta.conversations.start_stream(
            inputs=conversation_history,
            agent_id="ag_019a3501501b768a901bef17061c2a20",
        )

        # 5. 处理并打印模型的回复
        assistant_response = ""
        print("模型: ", end="")
        for chunk in response:
            # 假设 chunk 是一个有 content 属性的对象
            if hasattr(chunk, 'content'):
                print(chunk.content, end="", flush=True)
                assistant_response += chunk.content
        print()

        # 6. 将模型的回复也添加到对话历史中
        conversation_history.append({"role": "assistant", "content": assistant_response})

    except Exception as e:
        print(f"Error: {str(e)}")
        # 如果出错，可以选择移除最后一条用户消息，以便重试
        conversation_history.pop()